/*
 * main.cc
 *
 *  Created on: Dec 6, 2014
 *      Author: root
 */
/******************************************************************************************************************/
#include <stdio.h>
#include <stdlib.h>
#include <nids.h>
#include <dirent.h>
#include <string.h>
#include "./util/global.h"
#include "./knn/knn.h"
#include "main.h"
/******************************************************************************************************************/
int main(int argc, char* argv[]) {
//	printf("hello\n");
	if (!GlobalInit()) {
		perror("GlobalInit error!");
		exit(-1);
	}
	DIR* dir = NULL;
	dirent* ent = NULL;
	if (!(dir = opendir(PCAP_FILE_DIR))) {
		perror("opendir error!");
		return -1;
	}
	while ((ent = readdir(dir))) {
		// every pcap file related to one label.
		if (ent->d_type == DT_DIR || !strstr(ent->d_name, ".pcapng"))
			// not a file but a directory or not a .pcapng file.
			continue;
		char FQN[64] = { 0 };	// full qualified name
		sprintf(FQN, "%s%s", PCAP_FILE_DIR, ent->d_name);
		printf("%s\n", FQN);
		int len =
				strchr(ent->d_name, '.') ?
						strchr(ent->d_name, '-') - ent->d_name : -1;
		memset(g_label, 0, sizeof(g_label));
		strncpy(g_label, ent->d_name, len);
		printf("%s\n", g_label);
		nids_params.filename = FQN;
		if (!nids_init()) {
			fprintf(stderr, "%s\n", nids_errbuf);
			exit(-1);
		}
		nids_register_ip((void*) ip_func);
		nids_run();
		// run every pcap file and get KNN model.
		nids_exit();
	}
	free(ent);
	closedir(dir);
	/* live capture ip pkts. */
	printf("*************Start to Classify**************\n");
	g_tcp_m.clear();
	g_udp_m.clear();
	g_flag = true;
	nids_params.filename = (char*) PCAP_FILE_UDP;
	if (!nids_init()) {
		fprintf(stderr, "%s\n", nids_errbuf);
		exit(-1);
	}
	nids_register_ip((void*) ip_func);
	nids_run();
	nids_exit();
//	GlobalEnd();
	return 0;
}
/******************************************************************************************************************/
void ip_func(struct ip * a_packet) {
	if (g_flag) {
		SleepCertainTime();
	}
//	printf("%d ", ++g_ip_count);
	Ip* ip = (Ip*) a_packet;
//	ip->ToString();
	Socket src, dest;
	Socket client, server;
	src.addr = ip->saddr;
	dest.addr = ip->daddr;
//	struct timeval tmp = nids_last_pcap_header->ts;
	timeval tmp;
	gettimeofday(&tmp, NULL);
	double now = ConvertTimevalToDouble(tmp);
//	printf("%f\n", now);
	if (ip->IsTcp()) {
		Tcp* tcp = (Tcp*) ip->GetIpPayload();
		tcp->urg_ptr = ip->GetIpPayloadLen();
		src.port = tcp->source;
		dest.port = tcp->dest;
		if (tcp->IsSYN() && !tcp->IsACK() && !tcp->IsRST()) { // is syn and isn't ack, 1st handshake pkt;
				// client is src and server is dest
			client = src;
			server = dest;					// from c to s
			if (FindTcpStream(client, server)) {
				// find tcp stream, which means fin or rst do not recv and new tcp begins the same tuple4
				FreeTcpStream(client, server);
			}
			// new tcp stream
			Stream stream(client, server);
//			Statistic* stat = new Statistic();
			TcpStat* stat = new TcpStat();	// statistic data about this stream.
			stat->SetLabel(g_label);
			if (!g_tcp_m.insert(TCP_MAP::value_type(stream, stat)).second) {// tcp stream insert failed.
				perror("TCP_MAP insert error!");
				exit(-1);
			}
			stat->c_state = TCP_SYN_SENT;
			stat->s_state = TCP_SYN_RECV;
			// tcp stream insert success.
		} else if (tcp->IsSYN() && tcp->IsACK()) {// is syn and is ack, 2nd handshake pkt;
				// client is dest, server is src
			client = dest;
			server = src;					// from s to c
			TcpStat* stat = FindTcpStream(client, server);
			if (stat && stat->c_state == TCP_SYN_SENT
					&& stat->s_state == TCP_SYN_RECV) {
//				stat->attr[S_INIT_WIN] = ntohs(tcp->window);// server init window size
				stat->s_init_win = ntohs(tcp->window);
				stat->s_state = TCP_SYN_RECV;			// tcp state changed.
				stat->c_state = TCP_ESTABLISHED;
			}
		} else if (tcp->IsFIN() && tcp->IsACK()) {// is fin and is ack, close pkts
			if (!FreeTcpStream(src, dest)) {
				FreeTcpStream(dest, src);
			}
		} else if (tcp->IsRST()) {	// rst || rst ack
			if (!FreeTcpStream(src, dest)) {
				FreeTcpStream(dest, src);
			}
		} else if (tcp->IsACK()) {	// just ack or psh or urg, data pkts
			TcpStat* stat = FindTcpStream(src, dest);
			if (stat) {
				// src is client, dest is server
				client = src;
				server = dest;	// from c to s
				if (stat->s_state == TCP_SYN_RECV
						&& stat->c_state == TCP_ESTABLISHED) {// 3rd handshake pkt
//					stat->c_state = TCP_ESTABLISHED;// from now on we can send and recv pkts.
					stat->s_state = TCP_ESTABLISHED;
//					stat->attr[LAST_PKT_TIME] = now;
					stat->last_pkt_time = now;
				} else if (stat->c_state == TCP_ESTABLISHED
						&& stat->s_state == TCP_ESTABLISHED) {
					// normal data pkts
#if 0
					if (stat->cs_pkts == FIRST_N - 1) {
						// we get  FIRST_N pkts
						// classify

					}
#endif
					tcp->IsPSH() ? stat->cs_psh++ : 1;
					stat->cs_pkts++;
					stat->cs_bytes += tcp->GetTcpPayloadLen();
					stat->last_pkt_time = now;
				}
			}
			stat = FindTcpStream(dest, src);
			if (stat) {
				// dest is client, src is server;
				client = dest;
				server = src;								// from s to c
				if (stat->s_state == TCP_ESTABLISHED
						&& stat->c_state == TCP_ESTABLISHED) {
					// normal data pkt
					if (stat->sc_pkts == FIRST_N_TCP) {
						// we get  FIRST_N pkts
						// only when we get more than first_n pkts, we assume this stream is a p2p stream
						// otherwise we don't add it.
						if (g_flag) {
							// we need to classify unknown stream.
							KNN::Classify(stat);
							if (stat->label[0]) {
								printf("label: %s\n", stat->label);
							} else {
								printf("%s\n", "unknown");
							}
						} else {
							// not live capture, we need to build classify model.
							stat->ToString();
							AddToTcpVector(stat);
						}

					}
					tcp->IsPSH() ? stat->sc_psh++ : 1;
					stat->sc_pkts++;
					stat->sc_bytes += tcp->GetTcpPayloadLen();
					stat->last_pkt_time = now;
				}
			}
		}
//		tcp->ToString();
	} else if (ip->IsUdp()) {
		Udp* udp = (Udp*) ip->GetIpPayload();
		src.port = udp->source;
		dest.port = udp->dest;
//		udp->ToString();
		UdpStat* stat = FindUdpStream(src, dest);
		if (stat) {
			// we find a udp stream, and from c to s.
//			printf("%d ", udp->GetUdpPayloadLen());
			stat->cs_pkts++;
			stat->cs_bytes += udp->GetUdpPayloadLen();
			return;
		}
		stat = FindUdpStream(dest, src);
		if (stat) {
			// we find a udp stream, and from s to c
//			printf("%d ", udp->GetUdpPayloadLen());
			if (stat->sc_pkts == FIRST_N_UDP) {
				if (g_flag) {
					// we classify unknown stream
					KNN::Classify(stat);
					if (stat->label[0]) {
						printf("label: %s\n", stat->label);
					}
				} else {
					// we build udp stream already known model.
					stat->ToString();
					AddToUdpVector(stat);
				}
			}
			stat->sc_pkts++;
			stat->sc_bytes += udp->GetUdpPayloadLen();
			return;
		}
		// we do not find a udp stream, then we new one
		// and we assume that src is client, dest is server
		Stream stream(src, dest);
		stat = new UdpStat();
		if (!g_udp_m.insert(UDP_MAP::value_type(stream, stat)).second) {// udp stream insert failed.
			fprintf(stderr, "UDP_MAP insert error!\n");
			exit(-1);
		}
		stat->SetLabel(g_label);
		// first udp pkt.
//		printf("%d\n", udp->GetUdpPayloadLen());
	}
}
/******************************************************************************************************************/

/******************************************************************************************************************/
